package com.gcpproxy.response;

import com.gcpproxy.model.ChatCompletionRequest;
import com.gcpproxy.model.GeminiResponse;
import com.gcpproxy.model.Message;
import com.google.genai.types.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.Base64;

/**
 * 响应格式化器
 * 将Gemini API响应转换为OpenAI兼容格式
 */
@Component
public class ResponseFormatter {

    private static final Logger logger = LoggerFactory.getLogger(ResponseFormatter.class);

    /**
     * 格式化Chat Completion响应
     */
    public GeminiResponse formatChatCompletionResponse(
            GenerateContentResponse apiResponse,
            String model, String requestId,
            ChatCompletionRequest originalRequest) {
        GeminiResponse result = new GeminiResponse()
                .setId("chatcmpl-" + requestId)
                .setObject("chat.completion")
                .setModel(model)
                .setCreated(System.currentTimeMillis() / 1000);

        List<GeminiResponse.Choice> choices = buildChoices(apiResponse);
        result.setChoices(choices);
        result.setUsage(buildUsage(apiResponse, choices, originalRequest));

        return result;
    }

    private List<GeminiResponse.Choice> buildChoices(GenerateContentResponse apiResponse) {
        List<GeminiResponse.Choice> choices = new ArrayList<>();
        
        var candidates = apiResponse.candidates();
        if (candidates.isEmpty()) {
            logger.warn("candidates 不存在");
            choices.add(createEmptyChoice(0));
            return choices;
        }
        
        var candidatesList = candidates.get();
        for (int i = 0; i < candidatesList.size(); i++) {
            var candidate = candidatesList.get(i);
            List<Map<String, Object>> parts = extractPartsFromCandidate(candidate);
            Message message = new Message().setRole("assistant");
            message.setContent(formatContent(parts));
            
            choices.add(new GeminiResponse.Choice()
                    .setIndex(i)
                    .setFinishReason("stop")
                    .setMessage(message));
        }
        
        return choices;
    }

    private String formatContent(List<Map<String, Object>> parts) {
        if (parts.isEmpty()) {
            return ensureJsonFence("");
        }
        
        // 只有一个部分
        Map<String, Object> part = parts.get(0);
        if ("image_url".equals(part.get("type"))) {
            // 只有图片，返回 Markdown 格式的图片字符串
            Map<String, Object> imageUrl = (Map<String, Object>) part.get("image_url");
            String url = imageUrl != null ? (String) imageUrl.get("url") : "";
            if (url != null && !url.isEmpty()) {
                return "![image](" + url + ")";
            }
            return "";
        } else if ("text".equals(part.get("type"))) {
            // 只有文本，返回文本字符串
            return ensureJsonFence((String) part.get("text"));
        }
        
        return ensureJsonFence("");
    }

    private GeminiResponse.Choice createEmptyChoice(int index) {
        Message message = new Message().setRole("assistant").setContent("");
        return new GeminiResponse.Choice()
                .setIndex(index)
                .setFinishReason("stop")
                .setMessage(message);
    }

    private GeminiResponse.Usage buildUsage(GenerateContentResponse apiResponse, 
                                           List<GeminiResponse.Choice> choices,
                                           ChatCompletionRequest originalRequest) {
        int inputTokens = apiResponse.usageMetadata()
                .flatMap(um -> um.promptTokenCount())
                .orElse(0);
        
        int outputTokens = apiResponse.usageMetadata()
                .flatMap(um -> um.candidatesTokenCount())
                .orElse(0);
        
        if (inputTokens == 0) {
            inputTokens = estimateTokens(originalRequest);
        }
        if (outputTokens == 0) {
            outputTokens = estimateOutputTokens(choices);
        }

        int totalTokens = inputTokens + outputTokens;
        return new GeminiResponse.Usage()
                .setPromptTokens(inputTokens)
                .setCompletionTokens(outputTokens)
                .setTotalTokens(totalTokens)
                .setInputTokens(inputTokens)
                .setOutputTokens(outputTokens);
    }


    /**
     * 提取候选中的所有部分（文本和图片）
     */
    private List<Map<String, Object>> extractPartsFromCandidate(Candidate candidate) {
        List<Map<String, Object>> result = new ArrayList<>();
        
        var content = candidate.content();
        if (content.isEmpty()) {
            return result;
        }
        
        var parts = content.get().parts();
        if (parts.isEmpty()) {
            return result;
        }
        
        for (Part part : parts.get()) {
            if (part.text().isPresent()) {
                result.add(createTextPart(part.text().get()));
            } else if (part.inlineData().isPresent()) {
                var blob = part.inlineData().get();
                if (blob.mimeType().isPresent() && blob.data().isPresent()) {
                    result.add(createImagePart(blob.mimeType().get(), blob.data().get()));
                }
            }
        }
        
        return result;
    }

    private Map<String, Object> createTextPart(String text) {
        Map<String, Object> part = new HashMap<>();
        part.put("type", "text");
        part.put("text", text);
        return part;
    }

    private Map<String, Object> createImagePart(String mimeType, byte[] dataBytes) {
        String data = Base64.getEncoder().encodeToString(dataBytes);
        Map<String, Object> imageUrl = new HashMap<>();
        imageUrl.put("url", "data:" + mimeType + ";base64," + data);
        
        Map<String, Object> part = new HashMap<>();
        part.put("type", "image_url");
        part.put("image_url", imageUrl);
        return part;
    }

    private String ensureJsonFence(String text) {
        if (text == null || text.isEmpty()) return text;
        String trimmed = text.trim();
        // 如果已经有 ```json 或 ``` 包裹，则直接返回
        if (trimmed.startsWith("```json") || trimmed.startsWith("```") || trimmed.startsWith("{") || trimmed.startsWith("[")) {
            // 如果是纯 JSON（以 { 或 [ 开头），加上围栏
            if (trimmed.startsWith("{") || trimmed.startsWith("[")) {
                return "```json\n" + trimmed + "\n```";
            }
            return text;
        }
        return text;
    }

    private int estimateTokens(ChatCompletionRequest request) {
        if (request.getMessages() == null) return 0;

        int total = 0;
        for (var message : request.getMessages()) {
            total += estimateContentTokens(message.getContent());
        }
        return total;
    }

    private int estimateOutputTokens(List<GeminiResponse.Choice> choices) {
        int total = 0;
        for (GeminiResponse.Choice c : choices) {
            Message msg = c.getMessage();
            if (msg != null) {
                total += estimateContentTokens(msg.getContent());
            }
        }
        return total;
    }

    @SuppressWarnings("unchecked")
    private int estimateContentTokens(Object content) {
        if (content instanceof String) {
            return ((String) content).length() / 4;
        }
        if (content instanceof List) {
            int total = 0;
            for (Object item : (List<?>) content) {
                if (item instanceof com.gcpproxy.model.ContentItem) {
                    com.gcpproxy.model.ContentItem contentItem = (com.gcpproxy.model.ContentItem) item;
                    if (contentItem.isText() && contentItem.getText() != null) {
                        total += contentItem.getText().length() / 4;
                    } else if (contentItem.isImage()) {
                        total += 85;
                    }
                } else if (item instanceof Map) {
                    Map<String, Object> map = (Map<String, Object>) item;
                    if ("text".equals(map.get("type"))) {
                        String text = (String) map.get("text");
                        if (text != null) {
                            total += text.length() / 4;
                        }
                    } else if ("image_url".equals(map.get("type"))) {
                        total += 85;
                    }
                }
            }
            return total;
        }
        return 0;
    }

}